#include <limits.h>
#include <math.h>
#include <pthread.h>
#include <semaphore.h>
#include <sys/time.h>

#include "protocol.h"
#include "storage.h"

#define ALPHA 0.125
#define BETA 0.4

uint32_t crc_table[256];
sem_t enroute;
int remote_socket, local_socket;
struct sockaddr_in remote_name, local_name;
struct buffer send_buffer, recv_buffer;
time_t rto = 50;
time_t srtt = 50;
time_t rttvar = 25;

uint32_t crc32(uint8_t *data, unsigned int length) {
    uint32_t i, crc = 0xFFFFFFFF;
    for (i = 0; i < length; i++)
        crc = ((crc >> 8) & 0x00FFFFFF) ^ crc_table[(crc ^ *data++) & 0xFF];
    return (crc ^ 0xFFFFFFFF);
}

void gen_table() {
    uint32_t poly = 0xEDB88320L;
    int i, j;
    for (i = 0; i < 256; i++) {
        uint32_t crc = i;
        for (j = 8; j > 0; j--) {
            if (crc & 1) {
                crc = (crc >> 1) ^ poly;
            } else {
                crc >>= 1;
            }
        }
        crc_table[i] = crc;
    }
}

time_t time_ago(struct timeval *t) {
    struct timeval now;
    gettimeofday(&now, NULL);
    return (now.tv_sec - t->tv_sec)*1000 + (now.tv_usec - t->tv_usec)/1000;
}

void *local_to_remote() {
    socklen_t namelen = sizeof(local_name);
    uint32_t sequence = 0;
    struct sockaddr_in new_dest;
    
    while (1) {
        ssize_t amt;
        tcpd_container message;
        amt = recvfrom(local_socket, &message, sizeof message, 0, (struct sockaddr *)&local_name, &namelen);
        if (amt < 0)
            perror("local_to_remote recvfrom");
        new_dest.sin_family = AF_INET;
        new_dest.sin_port = TROLL_PORT;
        new_dest.sin_addr.s_addr = INADDR_ANY;
        message.troll = message.dest;
        message.troll.sin_port = TCPD_REMOTE;
        message.seq = htonl(sequence);
        message.checksum = htonl(crc32((uint8_t *)&message + sizeof(struct sockaddr_in), amt - sizeof(struct sockaddr_in)));
        sem_wait(&enroute);
#ifdef DEBUG
        if (message.flags == FIN)
            printf("tcpd: End of data, closing connection\n");
#endif
        amt = sendto(remote_socket, &message, amt, 0, (struct sockaddr *)&new_dest, sizeof(local_name));
        put(&send_buffer, &message, amt);
        sequence += (amt - NET_HEADER);
#ifdef DEBUG
        printf("tcpd: sending %d bytes, SEQ = %d, cksum = %x\n", (int)amt, ntohl(message.seq), message.checksum);
#endif
    }
    pthread_exit(NULL);
    return 0;
}

void *remote_to_local() {
    socklen_t namelen = sizeof(local_name);
    uint32_t acknowledgement = 0;

    while (1) {
        ssize_t amt;
        tcpd_container message;
        uint32_t original_cksum;
        amt = recvfrom(remote_socket, &message, sizeof message, 0, (struct sockaddr *)&remote_name, &namelen);
        original_cksum = message.checksum;
        message.checksum = 0;
        message.checksum = htonl(crc32((uint8_t *)&message + sizeof(struct sockaddr_in), amt - sizeof(struct sockaddr_in)));
#ifdef DEBUG
        printf("tcpd: receiving %d bytes, SEQ = %d, cksum = %x\n", (int)amt, ntohl(message.seq), message.checksum);
#endif
        
        if (original_cksum != message.checksum) {
#ifdef DEBUG
            printf("tcpd: checksum %x is bad, expected %x\n", message.checksum, original_cksum);
#endif
        } else if (message.flags == FIN) {
            /* Handle the other side closing the connection replying with
             * FIN/ACK. we should wait for an ACK or retransmitted FIN, or time
             * out trying
             */
            message.flags = FIN | ACK;
            message.checksum = htonl(crc32((uint8_t *)&message + sizeof(struct sockaddr_in), NET_HEADER - sizeof(struct sockaddr_in)));
#ifdef DEBUG
            printf("tcpd: Received connection teardown message.\n");
#endif
        } else if (message.flags == (FIN | ACK)) {
            message.flags = ACK;
            message.checksum = htonl(crc32((uint8_t *)&message + sizeof(struct sockaddr_in), NET_HEADER - sizeof(struct sockaddr_in)));
            sendto(remote_socket, &message, NET_HEADER, 0, (struct sockaddr *)&message.dest, sizeof(local_name));
            rto = 50;
            srtt = 50;
            rttvar = 25;
            acknowledgement = 0;
        } else if (message.flags == ACK) {
            tcpd_container discard;
            size_t length;
            struct timeval original = get(&send_buffer, message.seq, &discard, &length);
            sem_post(&enroute);
            if (length > 0) {
                time_t r = time_ago(&original);
                rttvar = (1 - BETA) * rttvar + BETA * abs(srtt - r);
                srtt = (1 - ALPHA) * srtt + ALPHA * r;
                rto = srtt + 4*rttvar;
                rto = rto < 20 ? 20 : rto;
                rto = rto > 1000 ? 1000 : rto;
            }
        } else {
            message.flags = ACK;
            message.checksum = 0;
            message.checksum = htonl(crc32((uint8_t *)&message + sizeof(struct sockaddr_in), NET_HEADER - sizeof(struct sockaddr_in)));
#ifdef DEBUG
            printf("tcpd: ACK'ing with %x\n", message.checksum);
#endif
            if (sendto(remote_socket, &message, NET_HEADER, 0, (struct sockaddr *)&remote_name, sizeof(local_name)) < 0)
                perror(NULL);
            if (ntohl(message.seq) == acknowledgement) {
                amt = sendto(local_socket, &message, amt, 0, (struct sockaddr *)&message.dest, sizeof(local_name));
                acknowledgement += (amt - NET_HEADER);
                while (lowest_seq(&recv_buffer) == acknowledgement) {
                    tcpd_container next;
                    size_t length;
                    get(&recv_buffer, acknowledgement, &next, &length);
                    amt = sendto(local_socket, &next, length, 0, (struct sockaddr *)&message.dest, sizeof(local_name));
                    acknowledgement += (amt - NET_HEADER);
                }
            } else if (ntohl(message.seq) > acknowledgement) {
                put(&recv_buffer, &message, amt);
            }
        }
    }
    pthread_exit(NULL);
    return 0;
}

int main() {
    pthread_t threads[2];
    if (sem_init(&enroute, 0, 20))
        perror(NULL);
    gen_table();
    storage_init(&send_buffer);
    storage_init(&recv_buffer);
    
    remote_socket = socket(AF_INET, SOCK_DGRAM, 0);
    if (remote_socket < 0) {
        perror("socket");
        exit(EXIT_FAILURE);
    }
    remote_name.sin_family = AF_INET;
    remote_name.sin_port = TCPD_REMOTE;
    remote_name.sin_addr.s_addr = INADDR_ANY;
    if (bind(remote_socket, (struct sockaddr *)&remote_name, sizeof(remote_name))) {
        perror("bind");
        exit(EXIT_FAILURE);
    }
    
    local_socket = socket(AF_INET, SOCK_DGRAM, 0);
    if (local_socket < 0) {
        perror("socket");
        exit(EXIT_FAILURE);
    }
    local_name.sin_family = AF_INET;
    local_name.sin_port = TCPD_LOCAL;
    local_name.sin_addr.s_addr = INADDR_ANY;
    if (bind(local_socket, (struct sockaddr *)&local_name, sizeof(local_name))) {
        perror("bind");
        exit(EXIT_FAILURE);
    }
    
    if (pthread_create(&threads[0], NULL, local_to_remote, NULL)) {
        perror(NULL);
        exit(EXIT_FAILURE);
    }
    if (pthread_create(&threads[1], NULL, remote_to_local, NULL)) {
        perror(NULL);
        exit(EXIT_FAILURE);
    }
    
    while (1) {
        uint32_t seq;
        time_t timer;
        if (pthread_mutex_lock(&send_buffer.mutex))
            perror(NULL);
        if (send_buffer.tail != send_buffer.head) {
            tcpd_container resend;
            size_t length;
            seq = send_buffer.contents[send_buffer.tail].data.seq;
            timer = time_ago(&send_buffer.contents[send_buffer.tail].timestamp);
            if (pthread_mutex_unlock(&send_buffer.mutex))
                perror(NULL);
            if (timer < rto)
                usleep(rto * 1000);
            get(&send_buffer, seq, &resend, &length);
            if (length != 0) {
                struct sockaddr_in new_dest;
                new_dest.sin_family = AF_INET;
                new_dest.sin_port = TROLL_PORT;
                new_dest.sin_addr.s_addr = INADDR_ANY;
#ifdef DEBUG
                printf("tcpd: timed out, resending SEQ = %d, cksum = %x\n", ntohl(resend.seq), resend.checksum);
#endif
                rto *= 2;
                rto = rto > 1000 ? 1000 : rto;
                sendto(remote_socket, &resend, length, 0, (struct sockaddr *)&new_dest, sizeof(local_name));
                put(&send_buffer, &resend, length);
            }
        } else {
            if (pthread_mutex_unlock(&send_buffer.mutex))
                perror(NULL);
            usleep(500 * 1000);
        }
    }

    pthread_join(threads[0], NULL);
    pthread_join(threads[1], NULL);
    close(local_socket);
    close(remote_socket);
    return EXIT_SUCCESS;
}
